import os
import pandas as pd
from langchain_text_splitters import RecursiveCharacterTextSplitter
import xml.etree.ElementTree as ET


def traverse_dir(path, extension, files=None, files_name=None):
    if files_name is None:
        files_name = []
    if files is None:
        files = []
    for file in os.listdir(path):
        file_path = os.path.join(path, file)
        if os.path.isdir(file_path):
            traverse_dir(file_path, extension, files, files_name)
        else:
            tail = file_path.split('.')[-1]
            if tail in extension and not os.path.basename(file_path).startswith('~$'):
                files.append(file_path)
                files_name.append(os.path.basename(file_path))
    return zip(files, files_name)


def get_doc_xml(unziped_path):
    # unziped_path下包含了OFD.xml文件，我们解析该文件获取Document.xml文件
    ofd_path = os.path.join(unziped_path, 'OFD.xml')
    tree = ET.parse(ofd_path)
    # 获取根元素
    root = tree.getroot()
    # OFD.xml文件中可能包含多个Doc_i文件夹，document_xml_path里面存放所有Doc_i/Document.xml的路径
    document_xml_path = []
    doc_i_path = []
    # 遍历所有元素
    for elem in root.iter():
        # DocRoot中存放Document.xml文件信息
        if 'DocRoot' in elem.tag:
            document_xml_path.append(os.path.join(unziped_path, elem.text))
            ly_dirs = elem.text.split('/')
            doc_i_parent = '/'.join(ly_dirs[0:-1])
            doc_i_path.append(os.path.join(unziped_path, doc_i_parent))
    return document_xml_path, doc_i_path


def get_attaches(document_xml_path, doc_i_path):
    page_path = []
    attach_path = set()
    for i, document_path in enumerate(document_xml_path):
        tree = ET.parse(document_path)
        root = tree.getroot()
        for elem in root.iter():
            if elem.attrib.__contains__('BaseLoc'):
                page_path.append(os.path.join(doc_i_path[i], elem.attrib['BaseLoc']))
            if 'Attachments' in elem.tag:
                attach_path.add(os.path.join(doc_i_path[i], elem.text))
    return page_path, attach_path


def rename_attach_files(attach_path):
    attach_names = []
    attach_ids = []
    attach_files_dir = []
    for document_path in attach_path:
        attach_parent_path = document_path.replace(os.path.basename(document_path), '')
        tree = ET.parse(document_path)
        root = tree.getroot()
        for elem in root.iter():
            if elem.attrib.__contains__('Name'):
                attach_names.append(elem.attrib['Name'])
            else:
                if attach_names:
                    base_name = os.path.basename(elem.text)
                    attach_ids.append(base_name)
                    attach_files_dir.append(os.path.join(attach_parent_path, elem.text.replace(base_name, '')))
    if len(attach_names) == len(attach_ids):
        for i in range(len(attach_names)):
            key = os.path.join(attach_files_dir[i], attach_ids[i])
            value = os.path.join(attach_files_dir[i], attach_names[i])
            if not os.path.exists(value):
                os.rename(key, value)
    else:
        print('ofd的Attach文件重命名失败！')


def process_exception(unziped_path):
    """
    :param unziped_path: 解压后的ofd文件夹
    :return: 获取ofd文字内容的xml路径
    获取路径：OFD.xml -> Doc_i/Document.xml -> Pages/Page_i/Content.xml
    """
    document_xml_path, doc_i_path = get_doc_xml(unziped_path)
    # 通过解析Document.xml来获取到page内容，page中包含了ofd文本内容的路径信息
    page_path, attach_path = get_attaches(document_xml_path, doc_i_path)
    # 获取 ofd中的附件信息（如果有的话），并根据映射关系重命名附件
    rename_attach_files(attach_path)
    # 解析page中的内容，获取ofd的文本相关的内容
    ofd_content = get_content_xml(page_path)
    return ofd_content


def get_content_xml(iter_files):
    content_ls = []
    for fid, file_path in enumerate(iter_files):
        tree = ET.parse(file_path)
        # 获取根元素
        root = tree.getroot()
        # 遍历所有元素
        for elem in root.iter():
            text = elem.text
            if text is not None and len(text.strip()) > 0 and 'TextCode' in elem.tag:
                content_ls.append(text.strip())
    content = ' '.join(content_ls).strip()
    return content
